package com.algorithm.redblack;

/**
 * @Author: DOWN
 * @Date: 2022/8/1 16:50
 * @Comment: 红黑树
 */
public class RedBlackTree<K extends Comparable<K>, V> {
    public static final boolean RED = true;
    public static final boolean BLACK = false;

    private class Node {
        private K key;
        public V value;
        public Node left, right;
        public boolean color;

        public Node(K key, V value) {
            this.key = key;
            this.value = value;
            left = null;
            right = null;
            color = RED;
        }
    }

    private Node root;
    private int size;

    public RedBlackTree() {
        root = null;
        size = 0;
    }

    public int getSize() {
        return this.size;
    }

    public boolean isEmpty() {
        return this.size == 0;
    }

    private boolean isRed(Node node) {
        if (node == null) {
            return BLACK;
        }
        return node.color;
    }

    /**
     * 左旋转
     * @param node 根节点
     * @return  Node
     */
    private Node leftRotate(Node node) {
        Node x = node.right;

        node.right = x.left;
        x.left = node;

        x.color = node.color;
        node.color = RED;

        return x;
    }

    /**
     * 右旋转
     * @param node 根节点
     * @return Node
     */
    private Node rightRotate(Node node) {
        Node x = node.left;

        node.left = x.right;
        x.right = node;

        x.color = node.color;
        node.color = RED;

        return x;
    }

    /**
     * 颜色反转（红黑树的性质）
     * @param node 根节点
     */
    private void flipColors(Node node) {
        node.color = RED;
        node.left.color = BLACK;
        node.right.color = BLACK;
    }

    /**
     * 添加元素
     * @param key 位置值
     * @param value 元素值
     */
    public void add(K key, V value) {
        add(root, key, value);
        root.color = BLACK;
    }

    private Node add(Node node, K key, V value) {
        if (node == null) {
            size++;
            return new Node(key, value);
        }
        if (key.compareTo(node.key) < 0) {
            node.left = add(node.left, key, value);
        } else if (key.compareTo(node.key) > 0) {
            node.right = add(node.right, key, value);
        } else {
            node.value = value;
        }
        //左黑 右红
        if (isRed(node.right) && !isRed(node.left)) {
            node = leftRotate(node);
        }
        //左红 左左红
        if (isRed(node.left) && isRed(node.left.left)) {
            node = rightRotate(node);
        }
        //左红 右红
        if (isRed(node.left) && isRed(node.right)) {
            flipColors(node);
        }
        return node;
    }
}
